import networkx as nx
import matplotlib.pyplot as plt
import random
import json
from networkx.readwrite import json_graph
import matplotlib

# 修正中文显示问题
matplotlib.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 中文字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 处理负号问题

def generate_ba_graph(num_nodes, edges_per_new_node):
    """
    使用 Barabási–Albert 模型生成基础网络
    :param num_nodes: 总节点数
    :param edges_per_new_node: 每个新节点的边连接数量
    :return: 生成的图对象
    """
    G = nx.barabasi_albert_graph(num_nodes, edges_per_new_node)

    # 为每个节点添加 'type' 属性
    for node in G.nodes():
        G.nodes[node]['type'] = 'center' if len(list(G.neighbors(node))) > edges_per_new_node else 'leaf'

    return G

def apply_waxman_on_ba(G, alpha=0.5, beta=0.2):
    """
    在 BA 模型生成的基础上应用 Waxman 模型来增加连接
    :param G: 已经生成的 BA 模型图
    :param alpha: Waxman 模型的 alpha 参数
    :param beta: Waxman 模型的 beta 参数
    :return: 添加了 Waxman 边的图
    """
    # 获取所有节点对
    nodes = list(G.nodes())
    for i in range(len(nodes)):
        for j in range(i + 1, len(nodes)):
            u, v = nodes[i], nodes[j]
            if u != v and not G.has_edge(u, v):
                # 计算两个节点之间的距离（这里用随机值代替）
                distance = random.uniform(0, 1)
                # Waxman 模型的连接概率
                prob = alpha * (1 - distance) * (distance ** beta)
                if random.random() < prob:  # 根据概率添加边
                    G.add_edge(u, v)

    # 更新 'type' 属性
    for node in G.nodes():
        G.nodes[node]['type'] = 'center' if len(list(G.neighbors(node))) > 1 else 'leaf'

    return G

def visualize_topology(G, title="结合 BA 和 Waxman 模型的网络拓扑图"):
    """
    可视化网络拓扑
    :param G: 图对象
    :param title: 图的标题
    """
    pos = nx.spring_layout(G)  # 使用 spring 布局
    node_colors = ['red' if G.nodes[node]['type'] == 'center' else 'lightblue' for node in G.nodes()]
    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=500)
    nx.draw_networkx_edges(G, pos, edge_color='gray')
    nx.draw_networkx_labels(G, pos, font_size=10, font_color='black')
    plt.title(title)
    plt.show(block=True)  # 显示图形并保持窗口打开

def export_topology_data(G, file_name="ba_waxman_graph_data.json"):
    """
    导出拓扑数据到 JSON 文件
    :param G: 图对象
    :param file_name: 输出文件名
    """
    data = json_graph.node_link_data(G)
    with open(file_name, 'w') as f:
        json.dump(data, f, indent=4)
    print(f"拓扑数据已导出到 {file_name}")

def main():
    # 参数设置
    num_nodes = 50  # 总节点数
    edges_per_new_node = 2  # 每个新节点的边连接数量
    alpha = 0.1  # Waxman 模型的 alpha 参数
    beta = 0.1  # Waxman 模型的 beta 参数

    # 生成 BA 拓扑
    G = generate_ba_graph(num_nodes, edges_per_new_node)

    # 在 BA 拓扑上应用 Waxman 模型
    G = apply_waxman_on_ba(G, alpha, beta)

    # 可视化拓扑
    visualize_topology(G)

    # 导出数据
    export_topology_data(G, "ba_waxman_graph_data.json")

if __name__ == "__main__":
    main()
